import pandas as pd
import json
import numpy as np
from tqdm import tqdm
import ast
import random
def get_category_item_json():
    columns = ['user_id', 'item_id', 'category_id', 'behavior', 'timestamp']
    usecol = ['item_id', 'category_id']
    print("读取数据中")
    data = pd.read_csv("UserBehavior.csv", sep=',',names=columns, usecols=usecol)
    print("读取完毕")
    data = np.array(data)
    category_item_dic = {}
    item_idx_list = []
    category_idx_list = []
    for item_id, category_id in tqdm(data):
        item_id = int(item_id)
        category_id = int(category_id)
        item_idx_list.append(item_id)
        category_idx_list.append(category_id)
        if category_id not in category_item_dic.keys():
            category_item_dic[category_id] = [item_id]
        else:
            category_item_dic[category_id].append(item_id)

    item_idx_list = sorted(list(set(item_idx_list)))
    category_idx_list = sorted(list(set(category_idx_list)))
    item_idx_dic = {
        str(item_idx):idx+1 for idx, item_idx in enumerate(item_idx_list)
    }
    category_idx_dic = {
        str(category_idx):idx+1 for idx, category_idx in enumerate(category_idx_list)
    }


    for key in category_item_dic.keys():
        category_item_dic[key] = list(set(category_item_dic[key]))

    json.dump(category_item_dic, open("category_item.json", "w"))
    json.dump(item_idx_dic, open("itemid2idx.json", "w"))
    json.dump(category_idx_dic, open("category2idx.json", "w"))

def split_dataset():
    print("读取数据")
    user_behavior = pd.read_csv("active_user_actions.csv", sep=',')
    user_behavior['item_id'] = user_behavior['item_id'].apply(ast.literal_eval)
    user_behavior['item_category'] = user_behavior['item_category'].apply(ast.literal_eval)
    user_behavior['timestamp'] = user_behavior['timestamp'].apply(ast.literal_eval)
    print("读取数据完成")
    # 打乱顺序（非常关键）
    user_behavior = user_behavior.sample(frac=1, random_state=42).reset_index(drop=True)

    # 分割成三部分：8:1:1
    # 分割成三部分：8:1:1
    train_size = int(0.8 * len(user_behavior))
    val_size = int(0.1 * len(user_behavior))

    df_train = user_behavior[:train_size]
    df_val = user_behavior[train_size:train_size + val_size]
    df_test = user_behavior[train_size + val_size:]


    with open('category_item.json', 'r', encoding='utf-8') as f:
        category2item = json.load(f)

    def get_pos_neg_data(df):
        pos_data = pd.DataFrame(columns=['user_id', 'item_id','cat_id', 'inter_type', 'timestamp'])
        tmp_pos_data = pd.DataFrame(columns=['user_id', 'item_id','cat_id', 'inter_type', 'timestamp'])
        for i in range(0,3):
            tmp_pos_data['user_id'] = df['user_id']
            tmp_pos_data['item_id'] = df['item_id'].apply(lambda x: x[i])
            tmp_pos_data['cat_id'] = df['item_category'].apply(lambda x: x[i])
            tmp_pos_data['timestamp'] = df['timestamp'].apply(lambda x: x[i])
            tmp_pos_data['inter_type'] = 1
            pos_data = pd.concat([pos_data, tmp_pos_data], ignore_index=True)

        def get_neg_item(row, i):
            pos_cat = row['item_category'][i]
            behavior_item_ids = row['item_id']
            behavior_item_ids = [int(item_id) for item_id in behavior_item_ids]
            while 1:
                proc = random.random()
                if proc < 2:
                    cat = random.choice(list(category2item.keys()))
                else:
                    cat = pos_cat
                cat_items = category2item[str(cat)]
                for idx,item in enumerate(cat_items):
                    if int(item) not in behavior_item_ids:
                        return item, int(cat)

        neg_data = pd.DataFrame(columns=['user_id', 'item_id','cat_id', 'timestamp', 'inter_type'])
        tmp_neg_data = pd.DataFrame(columns=['user_id', 'item_id', 'cat_id', 'inter_type', 'timestamp'])
        for i in range(3):
            tmp_neg_data['user_id'] = df['user_id']
            tmp_neg_data['item_id'] = df.apply(lambda x: get_neg_item(x, i)[0],axis=1)
            tmp_neg_data['cat_id'] = df.apply(lambda x: get_neg_item(x, i)[1],axis=1)
            tmp_neg_data['timestamp'] = df['timestamp'].apply(lambda x: int(x[i]) + 10)
            tmp_neg_data['inter_type'] = 0
            neg_data = pd.concat([neg_data, tmp_neg_data], ignore_index=True)
        data = pd.concat([pos_data, neg_data], ignore_index=True)
        data = data.sort_values(['user_id', 'timestamp']).reset_index(drop=True)
        return data

    final_train_data = get_pos_neg_data(df_train)
    final_val_data = get_pos_neg_data(df_val)
    final_test_data = get_pos_neg_data(df_test)

    final_train_data.to_csv('./final_train_data2.csv')
    final_val_data.to_csv('./final_valid_data2.csv')
    final_test_data.to_csv('./final_test_data2.csv')





if __name__ == '__main__':
    # get_category_item_json()
    split_dataset()
